import java.util.*;

/**
 * Created with IntelliJ IDEA.
 * Description:
 * User: mac
 * Date: 2022-10-12
 * Time: 17:13
 */
public class TopK {

    /**
     * 求数组中前K个最小的元素
     * @param array
     * @param k
     * @return
     */
    public static int[] topk(int[] array ,int k){
        //1.创建一个大小为k的大根堆
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });
        //2.遍历数组中当中的元素，前k个元素放到队列中
        for (int i = 0; i < array.length; i++) {
            if (maxHeap.size() < k){
                maxHeap.offer(array[i]);
            }else {
                int top = maxHeap.peek();
                if (top > array[i]){
                    maxHeap.poll();
                    maxHeap.offer(array[i]);//将新元素放入其中
                }
            }
        }
        int[] ret = new int[k];
        for (int i = 0; i < k; i++) {
            ret[i] = maxHeap.poll();
        }
        return ret;
    }
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k){
        PriorityQueue<List<Integer>> maxHeap = new PriorityQueue<>(k, new Comparator<List<Integer>>() {
            @Override
            public int compare(List<Integer> o1, List<Integer> o2) {
                return o2.get(0) + o2.get(1) - (o1.get(0) + o1.get(1));
            }
        });

        for (int i = 0; i < Math.min(nums1.length, k); i++) {
            for (int j = 0; j < Math.min(nums2.length, k); j++) {
                if (maxHeap.size() < k){
                    List<Integer> list = new ArrayList<>();
                    list.add(nums1[i]);
                    list.add(nums2[j]);
                    maxHeap.offer(list);
                }else{
                    List<Integer> list = new ArrayList<>();
                    list.add(nums1[i]);
                    list.add(nums2[j]);
                    List<Integer> cnt = maxHeap.peek();
                    int top = cnt.get(0) + cnt.get(1);
                    if (top > nums1[i] + nums2[j]){
                        maxHeap.poll();
                        maxHeap.offer(list);
                    }
                }
            }
        }
        List<List<Integer>> ret = new ArrayList<>();
        for (int i = 0; i < k && !maxHeap.isEmpty(); i++) {
            ret.add(maxHeap.poll());
        }
        return ret;
    }
    public static void main(String[] args) {
        int[] array = {18,21,8,10,34,21};
        int[] ret = topk(array, 3);
        System.out.println(Arrays.toString(ret));
    }
}
